# 导入必要的库
import os  # 用于文件路径处理
from monai.transforms import RandCropd  # 导入MONAI的随机裁剪变换（3D适用）
from monai.transforms import RandSpatialCrop
from monai.transforms import CropForegroundd  # 新增导入：定位有效区域
from monai.transforms import Crop  # 导入MONAI的随机裁剪变换（3D适用） 
import numpy as np  # 用于数值计算
import torch  # PyTorch深度学习框架
from matplotlib import pyplot as plt  # 可视化工具（此处未直接使用，保留为潜在可视化准备）
from monai.data import (  # MONAI的数据加载和处理类
    CacheDataset,  # 内存缓存数据集（加速重复加载）
    DataLoader,  # 批处理数据加载器
    Dataset,  # 基础数据集类
    PersistentDataset,  # 磁盘缓存数据集（适用于大数据）
    load_decathlon_datalist,  # 加载Decathlon格式的数据集列表
    partition_dataset,  # 划分数据集（用于交叉验证）
    decollate_batch  # 拆分批量数据为单个样本
)
from monai.transforms import (  # MONAI的数据变换工具
    EnsureChannelFirstd,  # 确保通道维度在第一个位置（符合PyTorch格式）
    Compose,  # 组合多个变换为一个序列
    DeleteItemsd,  # 删除数据字典中的指定键值对
    LoadImaged,  # 加载图像（支持多种格式，此处用于医学图像）
    Orientationd,  # 统一图像方向（医学图像标准化）
    ScaleIntensityd,  # 图像强度归一化
    SpatialPadd,  # 空间维度填充（确保裁剪前尺寸足够）
    ToTensord  # 转换为PyTorch张量
)


class BrainDecathlonDataModule:
    """
    用于处理Decathlon格式Hippocampus数据集的类，适用于T1-T2模态风格迁移任务。
    功能：加载数据、预处理（裁剪、归一化等）、划分训练/验证/测试集、创建数据加载器。
    """

    def __init__(
            self,
            root_dir=".",  # 数据集根目录，默认当前目录
            fold=0,  # 交叉验证折数（0-3），用于划分训练集和验证集
            train_patch_size=(32, 32, 32),  # 训练时裁剪的3D补丁大小（深度、高度、宽度）
            num_samples=4,  # 每个图像生成的裁剪样本数（数据增强，增加训练多样性）
            batch_size=1,  # 批处理大小（每次加载的样本数）
            cache_rate=0.,  # 数据缓存到内存的比例（0=不缓存，1=全缓存，适用于小数据）
            cache_dir=None,  # 数据缓存到磁盘的目录（None=不使用磁盘缓存，适用于大数据）
            num_workers=4,  # 数据加载的进程数（加速数据读取）
            balance_sampling=True,  # （已移除实际逻辑，保留参数为兼容旧代码）
            train_transforms=None,  # 自定义训练数据变换（None则使用默认变换）
            val_transforms=None,  # 自定义验证数据变换（None则使用默认变换）
            **kwargs  # 其他可选参数（预留扩展）
    ):
        # 拼接数据集具体路径：Decathlon格式中Hippocampus数据通常放在Task04_Hippocampus子目录
        self.base_dir = root_dir + "/BraTS2021_T1T2/"
        self.fold = fold  # 保存交叉验证折数（用于后续划分数据集）
        self.batch_size = batch_size  # 保存批处理大小
        self.cache_dir = cache_dir  # 保存磁盘缓存目录
        self.cache_rate = cache_rate  # 保存内存缓存比例
        self.num_workers = num_workers  # 保存数据加载进程数
        print('================ Loading Hippocampus Dataset ================')  # 提示数据集加载开始


        # 定义训练集预处理流程（T1和T2图像同步变换，确保空间一致性）
        if train_transforms is None:  # 若未提供自定义变换，则使用默认变换
            self.train_transforms = Compose([  # Compose将多个变换组合为一个流水线
                # 加载T1和T2图像（从dataset.json中读取"t1"和"t2"字段的路径）
                # reader="NibabelReader"：使用Nibabel库加载医学图像（支持.nii格式）
                LoadImaged(keys=["t1", "t2"], reader="NibabelReader"),
                # 确保通道维度在前（PyTorch要求格式：[通道数, 深度, 高度, 宽度]）
                EnsureChannelFirstd(keys=["t1", "t2"]),
                # 统一图像方向为LPI（医学图像标准方向：Left-Posterior-Inferior）
                # 避免因扫描方向不同导致的空间不一致
                Orientationd(keys=["t1", "t2"], axcodes="LPI"),
                # 强度归一化到[0, 1]范围（消除不同设备扫描强度差异的影响）
                ScaleIntensityd(keys=["t1", "t2"], minv=0.0, maxv=1.0),
                        # 5. 关键新增：定位有效区域（非零脑区），后续裁剪仅在该区域内进行
                CropForegroundd(
                    keys=["t1", "t2"],
                    source_key="t1",  # 以T1的有效区域为基准（T1和T2有效区域一致）
                    margin=0,  # 有效区域外不保留边缘（可设5-10增加容错，避免裁到边缘）
                    allow_smaller=True  # 显式指定，消除警告
                ),
                # 若图像尺寸小于目标裁剪尺寸，用边缘值填充（确保后续能裁剪出指定大小的补丁）
                SpatialPadd(keys=["t1", "t2"], spatial_size=train_patch_size, mode="edge"),
                # 随机裁剪（T1和T2裁剪相同区域，保持模态间的空间对应关系）
                RandCropd(
                    keys=["t1", "t2"],  # 保持对t1和t2的裁剪
                    # 重点：通过cropper参数传入裁剪器实例，指定roi_size和是否随机大小
                    cropper=RandSpatialCrop(
                        roi_size=train_patch_size,  # 裁剪尺寸（与输入补丁大小一致）
                        random_size=False  # 固定裁剪尺寸（不随机调整大小）
                        ),
                    allow_missing_keys=False  # 可选：是否允许缺失keys，默认False
                ),
                # 删除不需要的元数据（如图像路径、 affine矩阵等，减少内存占用）
                DeleteItemsd(keys=["t1_meta_dict", "t2_meta_dict"]),
                # 转换为PyTorch张量（便于模型计算）
                ToTensord(keys=["t1", "t2"]),
                # 删除变换记录（MONAI自动添加的变换历史，非必需）
                DeleteItemsd(keys=["t1_transforms", "t2_transforms"])
            ])
        else:
            self.train_transforms = train_transforms  # 使用自定义训练变换


        # 定义验证集预处理流程（无随机操作，保持图像完整，用于评估模型真实性能）
        if val_transforms is None:  # 若未提供自定义变换，则使用默认变换
            self.val_transforms = Compose([
                # 加载T1和T2图像（同训练集）
                LoadImaged(keys=["t1", "t2"], reader="NibabelReader"),
                # 确保通道维度在前（同训练集）
                EnsureChannelFirstd(keys=["t1", "t2"]),
                # 统一方向为LPI（同训练集，确保与训练数据空间一致）
                Orientationd(keys=["t1", "t2"], axcodes="LPI"),
                # 强度归一化（同训练集，确保数值范围一致）
                ScaleIntensityd(keys=["t1", "t2"], minv=0.0, maxv=1.0),
                # 删除元数据（同训练集）
                DeleteItemsd(keys=["t1_meta_dict", "t2_meta_dict"]),
                # 转换为张量（同训练集）
                ToTensord(keys=["t1", "t2"]),
                # 删除变换记录（同训练集）
                DeleteItemsd(keys=["t1_transforms", "t2_transforms"]),
            ])
        else:
            self.val_transforms = val_transforms  # 使用自定义验证变换


    def _load_data_dicts(self, train=True):
        """
        加载数据集字典（包含图像路径等信息），并划分训练/验证集。
        参数：train - 若为True，加载训练集并划分训练/验证；若为False，加载测试集。
        返回：数据字典列表、数据集信息（样本名等）
        """
        if train:
            # 加载训练集数据字典：从dataset.json中读取"training"字段的内容
            # dataset.json是Decathlon格式的数据集描述文件，包含每个样本的t1和t2路径
            data_dicts = load_decathlon_datalist(
                os.path.join(self.base_dir, "dataset.json"),  # dataset.json的路径
                data_list_key="training",  # 读取json中的"training"字段（训练样本列表）
                base_dir=self.base_dir  # 图像路径的基准目录
            )
            # 打印总样本数，确认数据加载完整
            print(f"总训练样本数: {len(data_dicts)}")  # 新增：验证总样本量

            # 划分数据集为4折（交叉验证）：将数据分成4份，按fold选择其中1份作为验证集，其余为训练集
            data_dicts_list = partition_dataset(
                data_dicts,  # 原始数据字典列表
                num_partitions=4,  # 分成4份
                shuffle=False,  # 不打乱（保证交叉验证的稳定性）
                seed=0  # 随机种子（固定划分结果）
            )
            # 打印各折样本数，确认划分是否均匀
            print(f"各折样本数: {[len(d) for d in data_dicts_list]}")  # 新增：检查每折样本
            train_dicts, val_dicts = [], []  # 初始化训练集和验证集字典列表
            info = dict()  # 存储数据集信息（如样本名）
            # 关键修复：初始化训练集样本名列表（必须！否则extend会报错）
            info["train"] = []  
            info["val"] = []  # 可选：初始化验证集列表，避免后续KeyError
            
            # 遍历4折数据，划分训练集和验证集
            for i, data_dict in enumerate(data_dicts_list):
                if i == self.fold:  # 第fold折作为验证集
                    val_dicts.extend(data_dict)  # 将该折数据加入验证集
                    # 记录验证集样本名（从t1路径中提取，便于后续追踪）
                    info["val"] = [os.path.basename(d["t1"]).split('.')[0] for d in data_dict]  # 直接用当前折的data_dict
                else:  # 其余折作为训练集
                    train_dicts.extend(data_dict)  # 将该折数据加入训练集
                    # 累加当前训练折的样本名，而非覆盖
                    current_train_names = [os.path.basename(d["t1"]).split('.')[0] for d in data_dict]
                    info["train"].extend(current_train_names)

            # 打印前5个训练和验证样本名（确认数据划分正确）
            #print("Train samples (first 5): ", info["train"][:5])
            #print("Val samples (first 5): ", info["val"][:5])
            return train_dicts, val_dicts, info  # 返回训练集、验证集和信息

        else:
            # 加载测试集数据字典：从dataset.json中读取"validation"字段
            test_dicts = load_decathlon_datalist(
                os.path.join(self.base_dir, "dataset.json"),
                data_list_key="validation",  # 读取json中的"validation"字段（测试样本列表）
                base_dir=self.base_dir
            )
            return test_dicts, None, None  # 测试集无需划分，返回测试数据字典


    def setup(self, stage=None):
        """
        初始化数据集（根据不同阶段加载训练/验证/测试数据）。
        参数：stage - 阶段标识（"fit"=训练+验证，"validate"=仅验证，None=默认加载训练+验证）
        返回：数据集信息（样本名等）
        """
        if stage in (None, "fit"):  # 训练+验证阶段
            # 加载训练集和验证集数据字典
            train_data_dicts, val_data_dicts, info = self._load_data_dicts()

            # 根据缓存设置选择数据集类型
            if self.cache_rate is not None:  # 使用内存缓存（适用于数据量较小的情况）
                # 初始化训练集（应用训练变换，按cache_rate缓存到内存）
                self.trainset = CacheDataset(
                    data=train_data_dicts,  # 训练数据字典列表
                    transform=self.train_transforms,  # 训练预处理变换
                    cache_rate=self.cache_rate,  # 缓存比例（0=不缓存，1=全缓存）
                    num_workers=self.num_workers  # 缓存时使用的进程数
                )
                # 初始化验证集（应用验证变换）
                self.valset = CacheDataset(
                    data=val_data_dicts,
                    transform=self.val_transforms,
                    cache_rate=self.cache_rate,
                    num_workers=4  # 验证集缓存进程数固定为4
                )
            elif self.cache_dir is not None:  # 使用磁盘缓存（适用于数据量较大的情况）
                self.trainset = PersistentDataset(
                    data=train_data_dicts,
                    transform=self.train_transforms,
                    cache_dir=self.cache_dir  # 缓存到磁盘的目录
                )
                self.valset = PersistentDataset(
                    data=val_data_dicts,
                    transform=self.val_transforms,
                    cache_dir=self.cache_dir
                )
            else:  # 不使用缓存（基础数据集，每次加载都重新处理）
                self.trainset = Dataset(
                    data=train_data_dicts,
                    transform=self.train_transforms
                )
                self.valset = Dataset(
                    data=val_data_dicts,
                    transform=self.val_transforms
                )

        elif stage == "validate":  # 仅验证阶段
            # 加载验证集数据字典（复用_load_data_dicts方法）
            _, val_data_dicts, info = self._load_data_dicts()
            # 验证集全缓存到内存（加速验证过程）
            self.valset = CacheDataset(
                data=val_data_dicts,
                transform=self.val_transforms,
                cache_rate=1.0,  # 全缓存
                num_workers=4
            )

        return info  # 返回数据集信息


    def train_dataloader(self):
        """创建训练集数据加载器（用于模型训练时批量读取数据）"""
        return DataLoader(
            self.trainset,  # 训练集数据集
            batch_size=self.batch_size,  # 批处理大小
            pin_memory=True,  # 锁页内存（加速GPU数据传输）
            num_workers=self.num_workers,  # 数据加载进程数
            drop_last=True,  # 丢弃最后一个不完整的批次（避免批次大小不一致）
            shuffle=False  # 不打乱顺序（风格迁移任务中顺序对训练影响较小）
        )


    def val_dataloader(self):
        """创建验证集数据加载器（用于模型验证时批量读取数据）"""
        return DataLoader(
            self.valset,  # 验证集数据集
            batch_size=1,  # 验证时通常单样本处理（便于评估每个样本的细节）
            num_workers=4,  # 固定4个进程
            shuffle=False  # 不打乱（保证验证结果可复现）
        )


    def test_dataloader(self):
        """创建测试集数据加载器（用于模型测试时批量读取数据）"""
        # 加载测试集数据字典
        test_data_dicts, _, _ = self._load_data_dicts(train=False)
        # 初始化测试集（使用验证变换，保持与验证流程一致）
        testset = CacheDataset(
            data=test_data_dicts,
            transform=self.val_transforms,
            cache_rate=self.cache_rate,
            num_workers=self.num_workers
        )
        # 创建测试集数据加载器
        return DataLoader(
            testset,
            batch_size=1,  # 单样本处理
            num_workers=4,
            shuffle=False  # 不打乱
        )